package com.csx.governance.traffic.mq.rocketmq.hook;

import cn.hutool.core.util.ObjectUtil;
import com.csx.governance.traffic.core.TrafficGovernanceRuleLocator;
import com.csx.governance.traffic.core.TransportableAttributes;
import com.csx.governance.traffic.mq.rocketmq.model.RocketMqServiceRule;
import org.apache.rocketmq.client.hook.SendMessageContext;
import org.apache.rocketmq.client.hook.SendMessageHook;
import org.springframework.context.ApplicationContext;

import java.util.Map;


/**
 * 发送消息的钩子
 * @author cuisongxu
 * @date 2023/2/28 16:28
 */
public class TrafficProduceMessageHook implements SendMessageHook {

    private final TrafficGovernanceRuleLocator<RocketMqServiceRule> ruleTrafficGovernanceRuleLocator;

    private final String localApplicationName;

    public TrafficProduceMessageHook(TrafficGovernanceRuleLocator<RocketMqServiceRule> ruleTrafficGovernanceRuleLocator,
                                     ApplicationContext applicationContext) {
        this.ruleTrafficGovernanceRuleLocator = ruleTrafficGovernanceRuleLocator;
        this.localApplicationName = applicationContext.getEnvironment().getProperty("spring.application.name");
    }

    @Override
    public String hookName() {
        return "TrafficProduceMessageHook";
    }

    @Override
    public void sendMessageBefore(SendMessageContext context) {

        // 判断该生产者是否需要加入到灰度中
        if(!isNeedJoinGray()) {
            return;
        }

        // 将需要透传的属性放入到消息的UserProperties中, swimLaneId 就是在此处透传到消息中的
        Map<String, String> attributes = TransportableAttributes.getAttributes();
        attributes.forEach(context.getMessage()::putUserProperty);

        System.out.println("TrafficProduceMessageHook#sendMessageBefore");
    }

    @Override
    public void sendMessageAfter(SendMessageContext context) {
        System.out.println("TrafficProduceMessageHook#sendMessageAfter");
    }

    private boolean isNeedJoinGray() {
        RocketMqServiceRule rule = ruleTrafficGovernanceRuleLocator.locateRule(localApplicationName);
        return !ObjectUtil.isEmpty(rule) && rule.isEnabled();
    }
}
